# -*- coding: utf-8 -*-


from __future__ import print_function
import numpy as np

from envs import next_step_back_mp_coo_large_scale_twice_timeout, next_step_back_stable_mp_coo_large_scale_twice_timeout

from envs import next_step_back_mp_coo_large_scale_twice_timeout_noprofiling, next_step_back_stable_mp_coo_large_scale_twice_timeout_noprofiling



class Board_sparse(object):
    """board for the game"""

    def __init__(self, rank, world_size, **kwargs):
        
        self.width = int(kwargs.get('width', 32))
        self.height = int(kwargs.get('height', 3))
        self.n_in_row = int(kwargs.get('n_in_row', 3))
        
        # board states stored as a dict,
        # key: move as location on the board,
        # value: player as pieces type
        # self.states = {}
        
        # self.matrix_feature = matrix_feature
        # self.gflops = []
        # self.parameters = [512, 256, 2]
        # self.available_selection = [1, 1, 1]
        # need how many pieces in a row to win
       
        # self.players = [1, 2]  # player1 and player2
        # self.players = 1
        
        self.parameters_scale_cpu = np.array([(128, 4096, 128), (32, 1024, 32), (1, 32, 1)]) # (start, end, stride) ������
        
        self.rank = rank
        self.world_size = world_size
        
        self.matrix_mean_gflops = np.zeros(32, dtype = float)
        self.matrix_select_time = np.zeros(32, dtype = int)
        

    def init_board(self, sparse_name, matrix_feature, mean_gflops, std_gflops):
        # if self.width < self.n_in_row or self.height < self.n_in_row:
        #     raise Exception('board width and height can not be '
        #                     'less than {}'.format(self.n_in_row))
        # self.current_player = self.players[start_player]  # start player
        # keep available moves in a list
        # self.players = 1
        
        self.availables = list(range(self.height))
        
        self.sparse_name = sparse_name
        self.matrix_feature = matrix_feature
        # self.manual_feature = manual_feature
        self.gflops_list = []
        self.parameters = [512, 256, 2]
        self.available_selection = [1, 1, 1]
        
        # self.states = {}
        # self.last_move = -1
        # self.last_profiling, self.last_gflops = self.get_profiling()
        self.last_gflops = self.get_gflops()
        self.init_gflops = self.last_gflops
        
        self.mean_gflops = mean_gflops
        self.std_gflops = std_gflops
        
        # self.last_profiling = []
        self.last_move = -1
        
    def get_gflops(self,):

        block_num, thread_num, element_num = self.parameters[0], self.parameters[1], self.parameters[2]
        try:
            gflops = next_step_back_mp_coo_large_scale_twice_timeout_noprofiling(self.sparse_name, block_num, thread_num, element_num, self.rank, self.world_size)

        except Exception as e:
            
            flag_continue = True 
            while flag_continue:
                try:
                    gflops = next_step_back_stable_mp_coo_large_scale_twice_timeout_noprofiling(self.sparse_name, block_num, thread_num, element_num, self.rank, self.world_size) # get next state
   
                    flag_continue = False
                except:
                    print('large: ', self.sparse_name, " ", block_num, " ", thread_num, " ", element_num)
        
        return gflops
        
        
    def get_profiling(self,):

        block_num, thread_num, element_num = self.parameters[0], self.parameters[1], self.parameters[2]
        try:
            profiling_feature, gflops = next_step_back_mp_coo_large_scale_twice_timeout(self.sparse_name, block_num, thread_num, element_num, self.rank, self.world_size)
            assert(len(profiling_feature) != 0)
        except Exception as e:
            
            flag_continue = True 
            while flag_continue:
                try:
                    profiling_feature, gflops = next_step_back_stable_mp_coo_large_scale_twice_timeout(self.sparse_name, block_num, thread_num, element_num, self.rank, self.world_size) # get next state
                    assert(len(profiling_feature) != 0)
                    flag_continue = False
                except:
                    print('large: ', self.sparse_name, " ", block_num, " ", thread_num, " ", element_num)
        
        return profiling_feature, gflops
    
        # self.last_profiling = profling_feature
        # self.gflops.append(gflops)

    # def move_to_location(self, move):
    #     """
    #     3*3 board's moves like:
    #     6 7 8
    #     3 4 5
    #     0 1 2
    #     and move 5's location is (1,2)
    #     """
    #     h = move // self.width
    #     w = move % self.width
    #     return [h, w]

    # def location_to_move(self, location):
    #     if len(location) != 2:
    #         return -1
    #     h = location[0]
    #     w = location[1]
    #     move = h * self.width + w
    #     if move not in range(self.width * self.height):
    #         return -1
    #     return move

    def current_state(self):
        
        '''
        get profiling configure
        '''
        """return the board state from the perspective of the current player.
        state shape: 4*width*height
        """
        parameters_normalize = np.array([(self.parameters[0] - self.parameters_scale_cpu[0][0]) / (self.parameters_scale_cpu[0][1] - self.parameters_scale_cpu[0][0]), 
                                    (self.parameters[1] - self.parameters_scale_cpu[1][0]) / (self.parameters_scale_cpu[1][1] - self.parameters_scale_cpu[1][0]),
                                    (self.parameters[2] - self.parameters_scale_cpu[2][0]) / (self.parameters_scale_cpu[2][1] - self.parameters_scale_cpu[2][0])])
        
        
        # return (self.matrix_feature, np.concatenate((self.last_profiling, parameters_normalize, self.available_selection), axis=-1))
        return (self.matrix_feature, np.concatenate((parameters_normalize, self.available_selection), axis=-1))

        # square_state = np.zeros((4, self.width, self.height))
        # if self.states:
        #     moves, players = np.array(list(zip(*self.states.items())))
        #     move_curr = moves[players == self.current_player]
        #     move_oppo = moves[players != self.current_player]
        #     square_state[0][move_curr // self.width,
        #                     move_curr % self.height] = 1.0
        #     square_state[1][move_oppo // self.width,
        #                     move_oppo % self.height] = 1.0
        #     # indicate the last move location
        #     square_state[2][self.last_move // self.width,
        #                     self.last_move % self.height] = 1.0
        # if len(self.states) % 2 == 0:
        #     square_state[3][:, :] = 1.0  # indicate the colour to play
        # return square_state[:, ::-1, :]

        
        # square_state = np.zeros((4, self.width, self.height))
        # if self.states:
        #     moves, players = np.array(list(zip(*self.states.items())))
        #     move_curr = moves[players == self.current_player]
        #     move_oppo = moves[players != self.current_player]
        #     square_state[0][move_curr // self.width,
        #                     move_curr % self.height] = 1.0
        #     square_state[1][move_oppo // self.width,
        #                     move_oppo % self.height] = 1.0
        #     # indicate the last move location
        #     square_state[2][self.last_move // self.width,
        #                     self.last_move % self.height] = 1.0
        # if len(self.states) % 2 == 0:
        #     square_state[3][:, :] = 1.0  # indicate the colour to play
        
        # square_state = (self.matrix_feature, self.last_move)
            
        # return square_state

    def do_move(self, move, move_value):
        # self.states[move] = self.current_player
        # self.states[move] = self.players
        self.availables.remove(move)
        
        # self.players += 1
        # self.current_player = self.players
        # self.current_player = (
        #     self.players[0] if self.current_player == self.players[1]
        #     else self.players[1]
        # )
        self.last_move = move
        self.available_selection[move] = 0
        # print(move)
        # input()
        # self.parameters[move] = move_value 
        
        self.parameters[move] = move_value * self.parameters_scale_cpu[move][2] \
                                                        + self.parameters_scale_cpu[move][0] # get 
        
        # self.last_profiling, gflops = self.get_profiling()
        gflops = self.get_gflops()
        
        # self.gflops_list.append((gflops - self.last_gflops) / self.init_gflops)
        # self.gflops_list.append((gflops - self.last_gflops) / self.mean_gflops)
        
        
        # normalize_gflops = (gflops - self.mean_gflops) / self.std_gflops
        # last_normalize_gflops = (self.last_gflops - self.mean_gflops) / self.std_gflops
        # self.gflops_list.append(normalize_gflops - last_normalize_gflops)
        
        self.gflops_list.append((gflops - self.last_gflops) / self.std_gflops)
        

        self.last_gflops = gflops
        
        # print(self.gflops_list[-1])
        

    def has_a_winner(self):
        # width = self.width
        height = self.height
        # states = self.states
        # n = self.n_in_row

        moved = list(set(range(height)) - set(self.availables))
        if (len(moved) == 3):
            return True
        # if len(moved) < self.n_in_row *2-1:
        #     return False, -1

        # for m in moved:
        #     h = m // width
        #     w = m % width
        #     player = states[m]

        #     if (w in range(width - n + 1) and
        #             len(set(states.get(i, -1) for i in range(m, m + n))) == 1):
        #         return True, player

        #     if (h in range(height - n + 1) and
        #             len(set(states.get(i, -1) for i in range(m, m + n * width, width))) == 1):
        #         return True, player

        #     if (w in range(width - n + 1) and h in range(height - n + 1) and
        #             len(set(states.get(i, -1) for i in range(m, m + n * (width + 1), width + 1))) == 1):
        #         return True, player

        #     if (w in range(n - 1, width) and h in range(height - n + 1) and
        #             len(set(states.get(i, -1) for i in range(m, m + n * (width - 1), width - 1))) == 1):
        #         return True, player

        return False

    def game_end(self):
        """Check whether the game is ended or not"""
        win = self.has_a_winner()
        if win:
            return True
        elif not len(self.availables):
            return True
        return False

    # def get_current_player(self):
    #     return self.current_player


class Game_sparse(object):
    """game server"""

    def __init__(self, board, **kwargs):
        self.board = board

    # def graphic(self, board, player1, player2):
    #     """Draw the board and show game info"""
    #     width = board.width
    #     height = board.height

    #     print("Player", player1, "with X".rjust(3))
    #     print("Player", player2, "with O".rjust(3))
    #     print()
    #     for x in range(width):
    #         print("{0:8}".format(x), end='')
    #     print('\r\n')
    #     for i in range(height - 1, -1, -1):
    #         print("{0:4d}".format(i), end='')
    #         for j in range(width):
    #             loc = i * width + j
    #             p = board.states.get(loc, -1)
    #             if p == player1:
    #                 print('X'.center(8), end='')
    #             elif p == player2:
    #                 print('O'.center(8), end='')
    #             else:
    #                 print('_'.center(8), end='')
    #         print('\r\n\r\n')

    def start_play(self, player, sparse_name, matrix_feature, mean_gflops, std_gflops,):
        """start a game between two players"""

        self.board.init_board(sparse_name, matrix_feature, mean_gflops, std_gflops)
                   
        while True:
            move, move_value, move_prob, move_value_prob = player.get_action(self.board)
                                # no temp | no return prob
            # print(move, move_prob)
            # print(move_value, move_value_prob)
            # input()
            self.board.do_move(move, move_value)
            end = self.board.game_end()
            if end:
                return self.board.parameters, self.board.init_gflops, self.board.last_gflops


    def start_self_play(self, player, sparse_name, matrix_feature, mean_gflops, std_gflops, temp=1e-3):
        """ 
        start a self-play game using a MCTS player, 
        and reuse the search tree,
        and store the self-play data: (state, mcts_probs, z) 
        for training
        """
        self.board.init_board(sparse_name, matrix_feature, mean_gflops, std_gflops)
        # p1, p2 = self.board.players
        # states, mcts_probs, current_players = [], [], []
        
        states, mcts_probs, mcts_value_probs = [], [], []
        actual_move = []
        while True:
            move, move_probs, move_value, move_value_probs = player.get_action(self.board,
                                                            temp=temp,
                                                            return_prob=1)
            #here
            # print(move, move_value)
            # input()
            
            # store the data
            states.append(self.board.current_state())
            mcts_probs.append(move_probs)
            mcts_value_probs.append(move_value_probs)
            actual_move.append(move)
            
            # current_players.append(self.board.current_player)
            # perform a move
            self.board.do_move(move, move_value)
            
            # if is_shown:
            #     self.graphic(self.board, p1, p2)
            end = self.board.game_end()
            if end:
                future_reward_list = []
                future_reward = np.zeros((1))
                reward_list = self.board.gflops_list
                for i in reversed(range(len(reward_list))):
                    future_reward = future_reward + reward_list[i] ####
                    future_reward_list.append(future_reward)
                future_reward_list = list(reversed(future_reward_list))
                # print(self.board.parameters)
                # print(future_reward_list)
                # input()
                
                ## reset MCTS root node
                player.reset_sparse()
                
                priority_compare = (self.board.last_gflops - self.board.init_gflops) / self.board.std_gflops
                priority_compare_list = [priority_compare for x in range(len(reward_list))]
                
                return zip(states, mcts_probs, mcts_value_probs, future_reward_list, actual_move, priority_compare_list), \
                        self.board.parameters, self.board.init_gflops, self.board.last_gflops
                    
                # # winner from the perspective of the current player of each state
                # winners_z = np.zeros(len(current_players))
                # if winner != -1:
                #     winners_z[np.array(current_players) == winner] = 1.0
                #     winners_z[np.array(current_players) != winner] = -1.0
                
                # if is_shown:
                #     if winner != -1:
                #         print("Game end. Winner is player:", winner)
                #     else:
                #         print("Game end. Tie")
                        
                # return zip(states, mcts_probs, winners_z)

                # return winner, zip(states, mcts_probs, winners_z)

    # def start_self_play_mp(self, player, sparse_name, matrix_feature, temp=1e-3):
    #     """ 
    #     start a self-play game using a MCTS player, 
    #     and reuse the search tree,
    #     and store the self-play data: (state, mcts_probs, z) 
    #     for training
    #     """
    #     self.board.init_board(sparse_name, matrix_feature)
    #     # p1, p2 = self.board.players
    #     # states, mcts_probs, current_players = [], [], []
        
    #     states, mcts_probs, mcts_value_probs = [], [], []
    #     actual_move = []
    #     while True:
    #         move, move_probs, move_value, move_value_probs = player.get_action(self.board,
    #                                                             temp=temp,
    #                                                             return_prob=1)
    #         # input()
    #         # store the data
    #         states.append(self.board.current_state())
    #         mcts_probs.append(move_probs)
    #         mcts_value_probs.append(move_value_probs)
    #         actual_move.append(move)
            
    #         # current_players.append(self.board.current_player)
    #         # perform a move
    #         self.board.do_move(move, move_value)
            
    #         # if is_shown:
    #         #     self.graphic(self.board, p1, p2)
    #         end = self.board.game_end()
    #         if end:
    #             future_reward_list = []
    #             future_reward = np.zeros((1))
    #             reward_list = self.board.gflops_list
    #             for i in reversed(range(len(reward_list))):
    #                 future_reward = future_reward + (reward_list[i] / self.board.init_gflops)
    #                 future_reward_list.append(future_reward)
    #             future_reward_list = list(reversed(future_reward_list))
                
    #             ## reset MCTS root node
    #             player.reset_sparse()
                
    #             return states, mcts_probs, mcts_value_probs, future_reward_list, actual_move, self.board.parameters
                    
    #             # # winner from the perspective of the current player of each state
    #             # winners_z = np.zeros(len(current_players))
    #             # if winner != -1:
    #             #     winners_z[np.array(current_players) == winner] = 1.0
    #             #     winners_z[np.array(current_players) != winner] = -1.0
    #             # # reset MCTS root node
    #             # player.reset_player()
    #             # if is_shown:
    #             #     if winner != -1:
    #             #         print("Game end. Winner is player:", winner)
    #             #     else:
    #             #         print("Game end. Tie")
                        
    #             # return zip(states, mcts_probs, winners_z)

    #             # return winner, zip(states, mcts_probs, winners_z)
